from langchain_core.messages import (
    AIMessage,
    HumanMessage,
    SystemMessage,
    trim_messages # 修剪消息
)
from langchain.text_splitter import CharacterTextSplitter
from langchain_deepseek import ChatDeepSeek

messages = [
    SystemMessage("you're a good assistant, you always respond with a joke."),
    HumanMessage("i wonder why it's called langchain"),
    AIMessage(
        'Well, I guess they thought "WordRope" and "SentenceString" just didn\'t have the same ring to it!'
    ),
    HumanMessage("and who is harrison chasing anyways"),
    AIMessage(
        "Hmmm let me think.\n\nWhy, he's probably chasing after the last cup of coffee in the office!"
    ),
    HumanMessage("what do you call a speechless parrot"),
]

trimmed_messages = trim_messages(
    messages,
    max_tokens=45, # 最大token数
    token_counter=ChatDeepSeek(model="gpt-4o-mini"), # 使用gpt-4o-mini模型计算token
    strategy="last", # 修剪策略，保留最后一条消息, 也可以选择"first"
    allow_partial=True, # 允许部分修剪
    end_on=HumanMessage, # 修剪结束条件，当遇到HumanMessage时停止修剪
    start_on=AIMessage, # 修剪开始条件，当遇到AIMessage时开始修剪
    include_system=True, # 是否包含SystemMessage
    text_splitter=CharacterTextSplitter(
        chunk_size=10, # 每个chunk的大小
        chunk_overlap=2 # 每个chunk的overlap
        ), # 使用CharacterTextSplitter进行文本分割
    )

print(trimmed_messages)


